#!/usr/bin/env python3
"""
将单卡版DexiNed权重转换为多卡(DataParallel)可直接加载的格式。
"""

import argparse
import os
from collections import OrderedDict
from typing import Any, Tuple

import torch


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="DexiNed权重前缀转换辅助脚本")
    parser.add_argument("--input", required=True, help="原始权重文件路径")
    parser.add_argument("--output", required=True, help="转换后权重输出路径")
    parser.add_argument("--overwrite", action="store_true", help="允许覆盖已存在的输出文件")
    return parser.parse_args()


def _assert_file_exists(path: str) -> None:
    expanded = os.path.abspath(os.path.expanduser(path))
    assert os.path.isfile(expanded), f"权重文件不存在: {expanded}"


def _load_checkpoint(path: str) -> Any:
    checkpoint = torch.load(path, map_location="cpu")
    assert checkpoint is not None, f"无法加载权重: {path}"
    return checkpoint


def _locate_state_dict(obj: Any) -> Tuple[OrderedDict, str]:
    if isinstance(obj, OrderedDict):
        return obj, "root"
    if isinstance(obj, dict):
        if "state_dict" in obj:
            sd = obj["state_dict"]
            assert isinstance(sd, (dict, OrderedDict)), "state_dict 字段必须是映射"
            return OrderedDict(sd), "state_dict"
        if "model" in obj:
            sd = obj["model"]
            assert isinstance(sd, (dict, OrderedDict)), "model 字段必须是映射"
            return OrderedDict(sd), "model"
    raise AssertionError("未找到可转换的state_dict字段")


def _convert_keys(state_dict: OrderedDict) -> OrderedDict:
    converted = OrderedDict()
    for key, value in state_dict.items():
        new_key = key if key.startswith("module.") else f"module.{key}"
        converted[new_key] = value
    assert len(converted) == len(state_dict), "键数量不一致，转换失败"
    return converted


def _save_checkpoint(obj: Any, where: str, location: str, new_state_dict: OrderedDict) -> None:
    if location == "root":
        torch.save(new_state_dict, where)
        return
    assert isinstance(obj, dict), "原始checkpoint结构异常"
    obj[location] = new_state_dict
    torch.save(obj, where)


def main() -> None:
    args = parse_args()
    input_path = os.path.abspath(os.path.expanduser(args.input))
    output_path = os.path.abspath(os.path.expanduser(args.output))

    _assert_file_exists(input_path)

    assert args.overwrite or not os.path.exists(output_path), f"输出文件已存在: {output_path}"

    output_dir = os.path.dirname(output_path)
    if output_dir != "":
        os.makedirs(output_dir, exist_ok=True)

    checkpoint = _load_checkpoint(input_path)
    state_dict, location = _locate_state_dict(checkpoint)
    converted = _convert_keys(state_dict)

    assert len(converted) > 0, "转换结果为空"
    sample_key = next(iter(converted.keys()))
    assert sample_key.startswith("module."), "转换后的键前缀校验失败"

    _save_checkpoint(checkpoint, output_path, location, converted)
    print(f"已生成多卡兼容权重: {output_path}")


if __name__ == "__main__":
    main()
